import numpy as np 
import cv2 as cv
import os

max_faces = 15
ratio = 0.1 # the min ratio between the rect & image, pass the smaller
def filt_boxes(boxes, ratio):
    for box in boxes:
        box_w,box_h = box[2:4]
        if ((box_w < ratio) and (box_h < ratio)):
            return []
    return boxes

def parse(shape, label_file_path, src_img_dir):
    fr = open(label_file_path,'r')
    line = fr.readline().rstrip()
    imgs = []
    boxes = []
    while line:
        path = os.path.join(src_img_dir, line)
        img = cv.imread(path)
        h, w, c = img.shape
        img = cv.resize(img, shape)
        num = int(fr.readline().rstrip())
        boxes_list = []
        if(num <= max_faces):
            for n in range(num):
                box = np.asarray([float(x) for x in fr.readline().rstrip().split(" ")[:4]] + [0, 1]) # need to center
                box[:2] += box[2:4] / 2
                box[:4:2] /= w
                box[1:4:2] /= h
                boxes_list.append(box)
            # discard the one, too small
            boxes_list = filt_boxes(boxes_list, ratio)
            if len(boxes_list) > 0 : 
                imgs += [img]
                boxes += [np.asarray(boxes_list)]
        else:
            for n in range(num):
                fr.readline()
        line = fr.readline().rstrip()
        # if(len(imgs) >= 5000):
        #     break
    np.save("./class_dataset/train_img.npy", imgs)
    np.save("./class_dataset/train_face_pos.npy", boxes)

def parse_from_rocky(shape, label_file_path, src_img_dir):
    import json 
    labels_dict = json.load(open(label_file_path))

    # our ssd need xywh as input
    imgs = []
    boxes = []
    for label_dict in labels_dict:
        file_name = label_dict['filename']
        path = os.path.join(src_img_dir, file_name)
        img = cv.imread(path)
        h, w, c = img.shape
        if not (shape == img.shape[1:-1]): img = cv.resize(img, shape)     
        imgs += [img]
        boxes_list = []
        for box in label_dict['bboxes']:
            box = np.asarray(box + [0, 1], dtype='float32')
            box[2:4] = box[2:4] - box[:2]
            box[:2] += box[2:4] / 2
            box[:4:2] /= w
            box[1:4:2] /= h
            boxes_list += [box] 
        boxes += [np.asarray(boxes_list)]

    np.save("./class_dataset/train_img.npy", imgs)
    np.save("./class_dataset/train_face_pos.npy", boxes)

if __name__=="__main__":
    file_path = r"F:\Face-detection-with-mobilenet-ssd-master\dataset\out_train_multi/bboxes.json"
    source_img_dir = r"F:\Face-detection-with-mobilenet-ssd-master\dataset" 
    # # test
    imgs = np.load("./class_dataset/train_img.npy")
    poses = np.load("./class_dataset/train_face_pos.npy")
    scales = []
    for i, img in enumerate(imgs[:100]):
        h, w, c = img.shape
        for pos in poses[i]:
            pos[:4:2] *= w 
            pos[1:4:2] *= h
            pos[:2] -= pos[2:4] / 2
            pos[2:4] += pos[:2]
            x1, y1, x2, y2 = pos.astype("uint32")[:4]
            cv.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0))
            scales += [[(x2 - x1) / w, (y2 - y1) / h, i]]
        cv.imshow("img", img)
        cv.waitKey(5)

    scales = np.asarray(scales).ravel().reshape(-1,3)
    scales_min_w_h = np.argmin(scales[:-1], axis=0)
    scales_max_w_h = np.argmax(scales[:-1], axis=0)

    sorted_pos = np.argsort(scales[:-1], axis=0)
    mid_pos = sorted_pos[[len(sorted_pos) // 2 - 1, len(sorted_pos) // 2 + 1]]

    min_w_img = imgs[int(scales[scales_min_w_h[0]][2])]
    min_h_img = imgs[int(scales[scales_min_w_h[1]][2])]

    max_w_img = imgs[int(scales[scales_max_w_h[0]][2])]
    max_h_img = imgs[int(scales[scales_max_w_h[1]][2])]   

    img_list = [min_w_img, min_h_img, max_w_img, max_h_img]


    mid_w = [mid_pos[0][0], mid_pos[1][0]]
    mid_h = [mid_pos[0][1], mid_pos[1][1]]

    mid_w_img = imgs[scales[mid_w][:,2].astype("uint32")]
    mid_h_img = imgs[scales[mid_h][:,2].astype("uint32")]

    print(scales[scales_min_w_h[0]], scales[scales_max_w_h[1]], scales[mid_w], scales[mid_h])

    img_list += mid_w_img.tolist()
    img_list += mid_h_img.tolist()

    for img in img_list:
        cv.imshow("img", np.asarray(img, dtype="uint8"))
        cv.waitKey(0)

    parse_func = parse_from_rocky
    parse_func((64, 64), file_path, source_img_dir)
